package com.xhf;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;

public class DefineTransformer implements ClassFileTransformer {
    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
                            ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
// 如果当前加载的类是我们编写的测试类，进入修改。
        if (className.contains("cn/xhf/")) {
            try {
                // 从ClassPool获得CtClass对象
                final ClassPool classPool = ClassPool.getDefault();
                final CtClass clazz = classPool.get("com.xhf.App");

                //打印App类中的所有成员函数
                CtMethod[] methodList = clazz.getDeclaredMethods();
                for(CtMethod method: methodList){
                    System.out.println("premain method: "+ method.getName());
                }

                // 获取add函数并替换，$1表示函数的第一个入参
                CtMethod convertToAbbr = clazz.getDeclaredMethod("add");
                String methodBody = "{return $1 + $2 + 11;}";
                convertToAbbr.setBody(methodBody);

                // 在add函数体之前增加一段代码，同理也可以在函数尾部添加
                methodBody = "System.out.println(Integer.toString($1));";
                convertToAbbr.insertBefore(methodBody);

                // 返回字节码，并且detachCtClass对象
                byte[] byteCode = clazz.toBytecode();
                //detach的意思是将内存中曾经被javassist加载过的Date对象移除，如果下次有需要在内存中找不到会重新走javassist加载
                clazz.detach();
                return byteCode;
            } catch (Exception ex) {
                ex.printStackTrace();
            }
        }
        // 如果返回null则字节码不会被修改
        return null;

    }
}
